簡介

ONNX Runtime (ORT) 與其他推論框架相比，是個相對古老的框架。但是他的泛用性相當高，可以適用於幾乎任何模型上。而 ORT 不只專注在推論上，在訓練上也有相當程度的優化。且 ORT 支援多個平台，能用來部署在 Windows, Linux, Mac, Android, iOS 等平台上，功能與社群都相對成熟，也是個不錯的選擇。

可惜的是 ORT 在這波 LLM 的熱潮中並沒有太多鏡頭，但依然不妨礙筆者拿 ORT 蹭一篇文章，今天就來講講 ONNX 與 ONNX Runtime 的基本用法，並瞭解為何這個框架的缺點在哪，為何使得 ORT 並沒有受到 LLM 玩家們的青睞。

(Powered By Microsoft Designer)

概覽

ONNX

Open Neural Network Exchange, ONNX 是一種文件格式，專門用來儲存機器學習模型。因為用來做機器學習的框架很多，例如 Tensorflow, PyTorch 之類的，大家都有各自的模型格式。ONNX 希望可以將這些格式統合起來，讓來自不同框架的模型都能夠用相同的方式執行。

ONNX 主要由 Microsoft, Amazon, Meta 與 IBM 等排擠 Google 的小圈圈公司們一同制定規範與維護。因為 Google 不在其中，所以 Tensorflow 並沒有正式支援 ONNX 格式。雖然 TF 有些非官方支援，但看起來並不是很活躍的社群。

ONNX Runtime

ONNX Runtime 是用來運行 ONNX 格式模型的執行階段函式庫，支援多個平台、程式語言與硬體架構，並整合了多種硬體加速方法。在 ORT 裡面這些加速方法被稱為 Execution Provider (EP)，例如 CUDA 就是其中一種 EP。除了廣泛的整合以外，ORT 也有支援 8-Bit 量化的功能。

最近 ORT v1.16.0 宣佈支援 4-Bit 量化，但目前還沒有詳細的文件參考如何進行量化。

ORT 其中一個優點是對模型支援比較通用，多數的模型架構不需要特別調整，就能轉換成 ONNX 格式並執行。其他框架例如 llama.cpp 就是以 Llama 架構為主，若要支援其他模型架構，則要透過 llama.cpp 的後端框架 ggml 自己實做。

但這也算是 ORT 的缺點之一，為了提高框架通用性的緣故，與其他框架比起來比較不那麼 "Lightweight"，用起來也相對拖泥帶水一點。

如果要使用 Python 版操作 ORT 的話，可以透過 pip 安裝：

pip install onnx  # ONNX 格式支援
pip install onnxruntime      # For CPU
pip install onnxruntime-gpu  # For CUDA

選擇 CPU 或 GPU 其中一個安裝就好。

推論

Optimum

Optimum 是 Hugging Face 整合多種加速函式庫的套件，其中就有 ONNX Runtime 這個選項。透過 Optimum 我們可以輕鬆將 Transformer 模型轉換成 ONNX 格式與運行。

為了方便與快速，以下使用 OPT 125M 模型進行示範。大多數的操作都可以用在 Llama 2 7B 之類的模型，但轉換過程會相對耗時（轉換 Llama 7B 約需 15 分鐘），且吃的 RAM 也很多，需要謹慎操作。

首先安裝相關套件：

pip install optimum accelerate

透過 optimum-cli 指令將模型轉為 ONNX 格式：

optimum-cli export onnx \
    --model facebook/opt-125m \
    --task text-generation-with-past \
    opt-125m-onnx

Decoder LM 的部份需要使用 text-generation-with-past 來加上 Past Key Value (KV Cache) 的輸入節點，如果是 Encoder-Decoder 模型應該就不用。雖然這是支援 PKV 的版本，但筆者力有未逮，沒有研究出正確使用 PKV 的方法。因此這裡只示範無快取的推論方式，缺點是推論效率會隨著長度增長而越來越低。

我們可以透過 Optimum 整合好的 ORTModelForCausalLM 來讀取模型，而 Tokenizer 的部份則使用 HF Transformers 就好：

from optimum.onnxruntime import ORTModelForCausalLM as ModelCls
from transformers import GPT2TokenizerFast as TkCls

model_path = "opt-125m-onnx"
model = ModelCls.from_pretrained(model_path)
tk = TkCls.from_pretrained(model_path)

最簡單的用法是透過 pipeline 來使用：

from optimum.pipelines import pipeline

generate = pipeline(
    "text-generation",
    model=model,
    tokenizer=tk,
    accelerator="ort",
)

outputs = generate("Hello, ", max_new_tokens=32)

ORT

我們可以透過 ORT 進行比較底層的操作，首先使用 InferenceSession 讀取模型：

from onnxruntime import InferenceSession

sess = InferenceSession(
    "opt-125m-onnx/decoder_model.onnx",
    providers=["CUDAExecutionProvider"],
)

如果想使用 CPU 的話，可以將 CUDAExecutionProvider 改成 CPUExecutionProvider。

在進行推論之前，我們要先確認輸入節點的名稱是什麼，可以透過 sess.get_inputs() 取得相關資訊：

for node in sess.get_inputs():
    print(node)

除了節點的名稱以外，還會包含輸入的形狀大小可以參考。這一步確定了這個模型包含 input_ids 與 attention_mask 等兩個輸入，我們需要將輸入以 dict 的方式傳入：

from transformers import GPT2TokenizerFast as TkCls

tk: TkCls = TkCls.from_pretrained("opt-125m-onnx")
tokens = tk("Hello, ")

inputs = {
    "input_ids": [tokens["input_ids"]],
    "attention_mask": [tokens["attention_mask"]],
}

接著使用 sess.run() 進行推論：

outputs = sess.run(None, inputs)

這個輸出會是個 list，我們可以根據 sess.get_outputs() 來確定這些輸出分別是什麼：

for k, v in zip(sess.get_outputs(), outputs):
    print(k.name, v.shape)

"""
輸出結果：
logits          (1,  4, 50272)
present.0.key   (1, 12, 4, 64)
present.0.value (1, 12, 4, 64)
present.1.key   (1, 12, 4, 64)
present.1.value (1, 12, 4, 64)
...
"""

觀察輸出發現，除了 logits 以外剩下都是 KV Cache 的輸出。我們這邊簡單對 logits 取 argmax 做 Greedy Decode，並整理成一個迴圈：

# 固定推論 32 次
for _ in tqdm(range(32)):
    # 進行推論
    outputs = sess.run(None, inputs)

    # outputs[0] 為 logits
    # 使用 .argmax 做 Greedy Decode
    new_tokens = outputs[0].argmax(-1)[:, -1:]

    # 把新的 Token 與原本的輸入串接在一起
    inputs["input_ids"] = np.concatenate(
        (inputs["input_ids"], new_tokens),
        axis=1,
    )

    # Attention Mask 隨著輸入的長度變長
    inputs["attention_mask"] = np.concatenate(
        (inputs["attention_mask"], [[1]]),
        axis=1,
    )

# 將整份 input_ids 解碼出來即為完整的輸入輸出
print(tk.decode(inputs["input_ids"][0]))

以上程式碼整理在這份 Colab Demo 裡面。

失敗分享

以上 ORT 的推論是沒有使用到 KV Cache 的，理論上這會比較慢。但因為今天家裡在開烤肉大趴，筆者沒有成功研究出使用 KV Cache 推論的版本。在這裡分享大致的思路給各位參考，不過這個環節的程式碼運作是不正常的。

在轉換完的模型資料夾裡面有三個 .onnx 檔案：

opt-125m-onnx/decoder_model_merged.onnx
opt-125m-onnx/decoder_model.onnx
opt-125m-onnx/decoder_with_past_model.onnx

decoder_model.onnx 是沒有使用 KV Cache 的版本，而筆者測試了一下，應該是使用 decoder_model_merged.onnx 操作有 KV Cache 的模型。觀察其輸入節點，大致如下：

NodeArg(
    name='input_ids',
    type='tensor(int64)',
    shape=['batch_size', 'sequence_length']
)
NodeArg(
    name='attention_mask',
    type='tensor(int64)',
    shape=['batch_size', 'attention_mask_sequence_length']
)
NodeArg(
    name='past_key_values.0.key',
    type='tensor(float)',
    shape=['batch_size', 12, 'past_sequence_length', 64]
)
NodeArg(
    name='past_key_values.0.value',
    type='tensor(float)',
    shape=['batch_size', 12, 'past_sequence_length', 64]
)
...
NodeArg(name='use_cache_branch', type='tensor(bool)', shape=[1])

筆者猜測關鍵應該在 shape 裡面的 sequence_length, attention_mask_sequence_length, past_sequence_length 這三者上面，但是不太確定推論過程中他們分別應該是多少。

先將輸入 Tokens 放進字典裡面，但我不確定 use_cache_branch 要設定成什麼：

tokens = tk("Hello, ")

inputs = {
    "input_ids": [tokens["input_ids"]],
    "attention_mask": [tokens["attention_mask"]],
    "use_cache_branch": [False],
}

推論之前先給初始的 KV Cache，但我不確定第三個維度要給多大：

init_zero = np.zeros((1, 12, 0, 64), dtype=np.float32)
for i in range(12):
    inn_feed[f"past_key_values.{i}.key"] = init_zero
    inn_feed[f"past_key_values.{i}.value"] = init_zero

接下來用這個迴圈跑 Decode 流程：

outputs = sess.run(None, inputs)

out_tokens = tokens["input_ids"]
inputs["attention_mask"] = [[1]]
inputs["use_cache_branch"] = [False]
for _ in range(16):
    out_tokens.append(outputs[0].argmax(-1)[0][-1])
    inputs["input_ids"] = outputs[0].argmax(-1)[:, -1:]

    for k, v in zip(inn_names[2:], outputs[1:]):
        inputs[k] = np.concatenate((inputs[k], v), axis=2)
    print(inputs[k].shape)

    outputs = sess.run(None, inputs)

最後 Decode 出來的結果是 "</s>Hello, I have the the the the the the the the the the the the the the"，很明顯的不正確。不過筆者就研究到這裡，畢竟 ORT 也不是 LLM 很常用的框架。

最後是筆者嘗試對 OPT 做量化，但是也失敗了。指令參考官方文件提供：

optimum-cli onnxruntime quantize --avx512 \
    --onnx_model opt-125m-onnx \
    -o opt-125m-onnx-quant

但是最後轉換出來的模型，無論使用 Optimum Pipeline 還是 ORT 都無法執行。

結論

今天探索了古老的 ONNX Runtime 框架，並使用簡單的 OPT 125M 示範。過去筆者多將 ORT 應用在參數量 1B 以下的 Encoder 或 Encoder-Decoder 上，並將模型推論過程以 C 語言實做，可以放到 Android 裝置上運行，且速度非常快。對這些使用案例來說，採用 ORT 框架是相當方便的。

但如果實際應用到 Llama 2 7B 上，就能感受到 ONNX Runtime 在這個層級上有多不好用。轉換時間非常的久，中間產物消耗的硬碟空間與主記憶體都很大。最後就是推論的正確姿勢到底長怎樣，實在有點難研究出來，可參考的資源不是很多。

接下來將會來介紹有名的 llama.cpp 與其背後的 ggml 框架，祝大家中秋節快樂，明天見！

參考

• HF: Intel Llama 7B ONNX INT4
• GitHub: Optimum
• GitHub: Llama 2 ONNX
• GitHub: ONNX Script
• GitHub: ONNX Inference Examples
• Optimum Pipelines
• Microsoft Open Source Blog
• HF Blog: ONNX Runtime
• iThome News: ONNX Script